__author__ = 'guoshunw'

from collections import namedtuple
from struct import *

#FileHeader class definition
FileHeader = namedtuple('FileHeader', 'usMagic usNumSec ulTime ulSymbolOffset ulNumSymbol usOptHdrSZ usFlags')


def header2String(self):
    myList = list(self)
    from datetime import datetime

    myList[2:2] = [str(datetime.fromtimestamp(myList[2]))]
    retStr = ("usMagic\tusNumSec\tulTime\t\t\t\t\t\t\tulSymbolOffset\tulNumSymbol\tusOptHdrSZ\tusFlags\n"
              "%04X\t%04X\t\t%s(%08X)\t%08X\t\t%08X\t%04X\t\t%04X" % tuple(myList))
    return retStr + '\n' + '=' * 100 + '\n'

FileHeader.__str__ = header2String
FileHeader._unpack_format = 'HHLLLHH'

#OptionalHeader class definition
OptHeader = namedtuple('OptHeader',
    'usMagic usVersion ulTextSize ulInitDataSZ ulUninitDataSZ ulEntry ulTextBase ulDataBase')
OptHeader._unpack_format = 'HHLLLLLL'
OptHeader.__str__ = lambda self: (
    "usMagic usVersion ulTextSize ulInitDataSZ ulUninitDataSZ ulEntry ulTextBase ulDataBase\n"
    "%04X\t%04X\t\t%s(%08X)\t%08X\t\t%08X\t%04X\t\t%04X" % tuple(self))

SectionHeader = namedtuple('SectionHeader',
    'cName ulVSize ulVaddr ulSize ulSecOffset ulRelOffset ulLNOffset usNumRel usNumLN ulFlags')
SectionHeader._unpack_format = '8sLLLLLLHHL'
SectionHeader._toString_header = " cName\tulVSize\t\tulVaddr\t\tulSize\t\tulSecOffset\tulRelOffset\t" \
                                 "ulLNOffset\tusNumRel\tusNumLN\tulFlags\n"

def secHdr2String(self):
    myList = list(self)
    myList[0] = myList[0].rstrip('\0')
    return "\n%-08s%08X\t%08X\t%08X\t%08X\t%08X\t%08X\t%04X\t\t%04X\t%08X" % tuple(myList)


SectionHeader.__repr__ = secHdr2String

#Common Object File Format definition

class COFFFile:
    """This class represent a common object file format file which usually is .o or .obj generated by compiler. """

    def __init__(self, file):
        offset = 0
        file.seek(offset)
        # get buffer from COFF file
        buffer = file.read()

        self.fileHeader = FileHeader._make(unpack_from(FileHeader._unpack_format, buffer, offset))
        offset += calcsize(FileHeader._unpack_format)

        self.optHeader = None if self.fileHeader.usOptHdrSZ == 0 else\
        OptHeader._make(unpack_from(OptHeader._unpack_format, buffer, offset))
        offset += self.fileHeader.usOptHdrSZ

        self.sections = []
        secHdrSize = calcsize(SectionHeader._unpack_format)
        for i in range(self.fileHeader.usNumSec):
            secHdr = SectionHeader._make(unpack_from(SectionHeader._unpack_format, buffer, offset))
            self.sections.append(secHdr)
            offset += secHdrSize


    def __str__(self):
        optHeaderStr = '' if not self.optHeader else str(self.optHeader)
        return str(self.fileHeader) + optHeaderStr + SectionHeader._toString_header + str(self.sections)